#include <iostream>
#include <sstream>
#include <chrono>
#include "WritePng.h"

struct vec3{
    float x, y, z;
};

struct mat3{
    float v[3 * 3];
};

static void mat3MulVec3(mat3 mat, vec3 v, vec3 *outV){
    outV->x = mat.v[0] * v.x + mat.v[1] * v.y + mat.v[2] * v.z;
    outV->y = mat.v[3] * v.x + mat.v[4] * v.y + mat.v[5] * v.z;
    outV->z = mat.v[6] * v.x + mat.v[7] * v.y + mat.v[8] * v.z;
}

float clamp(float v, float min, float max){
    if(v < min){
        return min;
    }
    if(v > max){
        return max;
    }
    return v;
}

bool readFile(const char *filePath, size_t *outSize, uint8_t **outData){
    FILE *fp = fopen(filePath,"rb");
    if (fp){
        fseek(fp, 0, SEEK_END);
        *outSize = ftell(fp);
        *outData = static_cast<uint8_t *>(malloc(sizeof(uint8_t) * (*outSize)));
        fseek(fp, 0, SEEK_SET);
        fread(*outData, 1, *outSize, fp);
        fclose(fp);
        fp = nullptr;
        return true;
    }
    return false;
}

static float RGBYUV14075[256], RGBYUV03455[256], RGBYUV07169[256], RGBYUV17790[256];
static float YUVRGB02990[256], YUVRGB05870[256], YUVRGB01140[256], YUVRGB01690[256],
             YUVRGB03310[256], YUVRGB05000[256], YUVRGB04190[256], YUVRGB00810[256];
static void initQueryTable(){
    int i;
    for (i = 0; i < 256; i++) RGBYUV14075[i] = (float)1.4075 * (i - 128);
    for (i = 0; i < 256; i++) RGBYUV03455[i] = (float)0.3455 * (i - 128);
    for (i = 0; i < 256; i++) RGBYUV07169[i] = (float)0.7169 * (i - 128);
    for (i = 0; i < 256; i++) RGBYUV17790[i] = (float)1.7790 * (i - 128);

    for (i = 0; i < 256; i++) YUVRGB02990[i] = (float)0.2990 * i;
    for (i = 0; i < 256; i++) YUVRGB05870[i] = (float)0.5870 * i;
    for (i = 0; i < 256; i++) YUVRGB01140[i] = (float)0.1140 * i;
    for (i = 0; i < 256; i++) YUVRGB01690[i] = (float)0.1690 * i;
    for (i = 0; i < 256; i++) YUVRGB03310[i] = (float)0.3310 * i;
    for (i = 0; i < 256; i++) YUVRGB05000[i] = (float)0.5    * i;
    for (i = 0; i < 256; i++) YUVRGB04190[i] = (float)0.4190 * i;
    for (i = 0; i < 256; i++) YUVRGB00810[i] = (float)0.0810 * i;
}

struct YuvData{
    enum Type{
        NV12, NV21
    } type;
    uint32_t width;
    uint32_t height;
    size_t totalSize;
    uint32_t frameCount;
    uint8_t *yPlane;
    uint8_t *uvPlane;

    bool load(const char *filePath){
        //read yuv file
        if(!readFile(filePath, &totalSize, &yPlane)){
            std::cerr << "Could not read yuv file:" << filePath << std::endl;
            return false;
        }
        size_t yuvFrameSize = width * height * 3 / 2 * sizeof(uint8_t);
        uvPlane = yPlane + width * height;
        frameCount = totalSize % yuvFrameSize == 0 ? totalSize / yuvFrameSize : totalSize / yuvFrameSize + 1;
        return true;
    }
};

struct RgbData{
    uint32_t width;
    uint32_t height;
    size_t totalSize;
    uint32_t frameCount;
    uint8_t *data;

    bool load(const char *filePath){
        //read yuv file
        if(!readFile(filePath, &totalSize, &data)){
            std::cerr << "Could not read yuv file:" << filePath << std::endl;
            return false;
        }
        size_t yuvFrameSize = width * height * 3 * sizeof(uint8_t);
        //frameCount = totalSize % yuvFrameSize == 0 ? totalSize / yuvFrameSize : totalSize / yuvFrameSize + 1;
        frameCount = totalSize / yuvFrameSize;
        return true;
    }
};

static bool writeYuv(const char *filePath, size_t size, uint8_t *yuvBuffer){
    FILE *fp = fopen(filePath,"wb");
    if (fp){
            fwrite(yuvBuffer, size, 1, fp);
        fclose(fp);
        fp = nullptr;
        return true;
    } else {
        std::cerr << "Could not read rgb file:" << filePath << std::endl;
        return false;
    }
}

static bool writeRgb(const char *filePath, uint32_t width, uint32_t height, size_t size, uint8_t *rgbBuffer, bool isGenPng){
    FILE *fp = fopen(filePath,"wb");
    if (fp){
        if(isGenPng) {
            std::stringstream ss;
            ss << filePath << ".png";
            FILE *pngFp = fopen(ss.str().c_str(),"wb");
            if(pngFp){
                svpng(pngFp, width, height, rgbBuffer, 0);
                fclose(pngFp);
                pngFp = nullptr;
            }
        }
        fwrite(rgbBuffer, size, 1, fp);
        fclose(fp);
        fp = nullptr;
        return true;
    } else {
        std::cerr << "Could not read rgb file:" << filePath << std::endl;
        return false;
    }
}

bool YUV2RGB(const YuvData &yuvData, size_t rgbSize, uint8_t **rgbBuffer){
    if(rgbSize <= 0){
        return false;
    }
    uint8_t *rgbIndex = *rgbBuffer;
    uint8_t *yuvYIndex = yuvData.yPlane;
    uint64_t uvBufferOffset = 0;

    for(int i = 0; i < yuvData.frameCount; i++){
        for(int w = 0; w < yuvData.width; w++){
            for(int h = 0; h < yuvData.height; h++){
                int indexU, indexV;
                if(h % 2 == 0){
                    indexU = (w / 2 * yuvData.width + h) + uvBufferOffset;
                    indexV = indexU + 1;
                } else {
                    indexV = (w / 2 * yuvData.width + h) + uvBufferOffset;
                    indexU = indexV - 1;
                }

                bool nv21 = yuvData.type == YuvData::NV21;
                uint8_t Y = *yuvYIndex++;
                uint8_t U = yuvData.uvPlane[nv21 ? indexV : indexU];
                uint8_t V = yuvData.uvPlane[nv21 ? indexU : indexV];

//                float R = Y + 1.4075f * (V - 128);
//                float G = Y - 0.3455f * (U - 128) - 0.7169f * (V - 128);
//                float B = Y + 1.7790f * (U - 128);

                vec3 yuv = { (float)Y, (float)(U - 128), (float)(V - 128) };
                mat3 mat = {
                        1, 0, 1.4075f,
                        1, -0.3455f, -0.7169f,
                        1, 1.779f, 0
                };
                vec3 rgb{};
                mat3MulVec3(mat, yuv, &rgb);
                float R = rgb.x;
                float G = rgb.y;
                float B = rgb.z;

//                float R = Y + RGBYUV14075[V];
//                float G = Y - RGBYUV03455[U] - RGBYUV07169[V];
//                float B = Y + RGBYUV17790[U];

                R = clamp(R, 0x00, 0xFF);
                G = clamp(G, 0x00, 0xFF);
                B = clamp(B, 0x00, 0xFF);

                *rgbIndex++ = (uint8_t) R;
                *rgbIndex++ = (uint8_t) G;
                *rgbIndex++ = (uint8_t) B;
            }
        }
        uvBufferOffset += yuvData.width * yuvData.height * 3 / 2;
    }
    return true;
}

bool RGB2YUV(const RgbData &rgbData, YuvData::Type type, size_t yuvSize, uint8_t **yuvBuffer){
    if(yuvSize <= 0){
        return false;
    }
    uint8_t *yuvYIndex = *yuvBuffer;
    uint8_t *yuvUVIndex = *yuvBuffer + rgbData.width * rgbData.height;
    uint8_t *rgbIndex = rgbData.data;

    for(int i = 0; i < rgbData.frameCount; i++){
        for(int w = 0; w < rgbData.width; w++){
            for(int h = 0; h < rgbData.height; h++){
                uint8_t R = *rgbIndex++;
                uint8_t G = *rgbIndex++;
                uint8_t B = *rgbIndex++;

                float Y = YUVRGB02990[R] + YUVRGB05870[G] + YUVRGB01140[B];
                Y = clamp(Y, 0x00, 0xFF);
                *yuvYIndex++ = (uint8_t) Y;

                //sampler 2x2
                if(w % 2 == 0 && h % 2 == 0){
                    float U = (-YUVRGB01690[R] - YUVRGB03310[G] + YUVRGB05000[B]) + 128;
                    float V = (YUVRGB05000[R] - YUVRGB04190[G] - YUVRGB00810[B]) + 128;

                    U = clamp(U, 0x00, 0xFF);
                    V = clamp(V, 0x00, 0xFF);

                    bool isN21 = type == YuvData::NV21;
                    *yuvUVIndex++ = (uint8_t) (isN21 ? V : U);
                    *yuvUVIndex++ = (uint8_t) (isN21 ? U : V);
                }
            }
        }
    }
    return true;
}

int main(){
    const char *yuvFilePath = "rgb4_22_1600x1600_NV12.yuv";
    const char *yuvNewFilePath = "rgb4_22_1600x1600_NV12___new.yuv";
    const char *rgbFilePath = "rgb4_22_1600x1600_NV12.rgb";

//    const char *yuvFilePath = "/home/cxh/st/rgb2.yuv";
//    const char *yuvNewFilePath = "rgb2___new.yuv";
//    const char *rgbFilePath = "rgb2.rgb";

    const uint32_t imgWidth = 1600;
    const uint32_t imgHeight = 1600;
    const YuvData::Type yuvType = YuvData::NV21;

    // init query table
    initQueryTable();

    //test 1, yuv -> rgb
    {
        YuvData yuv{ .type = yuvType, .width = imgWidth, .height = imgHeight };
        yuv.load(yuvFilePath);

        auto startTime = std::chrono::steady_clock::now();
        size_t rgbSize = yuv.frameCount * yuv.width * yuv.height * 3 * sizeof(uint8_t);
        uint8_t *rgbBuffer = static_cast<uint8_t *>(malloc(rgbSize));
        if(!YUV2RGB(yuv, rgbSize, &rgbBuffer)){
            std::cerr << "Could not convert yuv to rgb" << std::endl;
            return -1;
        }
        auto enTime = std::chrono::steady_clock::now();
        printf("duration: %.4f\n", std::chrono::duration<float, std::milli>(enTime - startTime).count());
        writeRgb(rgbFilePath, yuv.width, yuv.height, rgbSize, rgbBuffer, true);
        free(rgbBuffer);
        free(yuv.yPlane);
    }

    //test2, rgb -> yuv
    {
        RgbData rgb{ .width = imgWidth, .height = imgHeight };
        if(!rgb.load(rgbFilePath)){
            return -1;
        }

        size_t yuvSize = rgb.frameCount * rgb.width * rgb.height * 3 / 2 * sizeof(uint8_t);
        uint8_t *yuvBuffer = static_cast<uint8_t*>(malloc(yuvSize));
        if(!RGB2YUV(rgb, YuvData::NV21, yuvSize, &yuvBuffer)){
            std::cerr << "Could not convert yuv to rgb" << std::endl;
            return -1;
        }
        writeYuv(yuvNewFilePath, yuvSize, yuvBuffer);

        free(yuvBuffer);
        free(rgb.data);
    }

    std::cout << "Convert end." << std::endl;

    return 0;
}